from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
    Trainer,
    TrainingArguments,
    DataCollatorWithPadding
)
from datasets import load_dataset
import evaluate
import torch

# 1. 加载数据集和分词器
dataset = load_dataset("imdb")
print(len(dataset))
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
print(tokenizer)
# 2. 数据预处理
def preprocess(examples):
    return tokenizer(examples["text"], truncation=True, max_length=512)

tokenized_data = dataset.map(preprocess, batched=True)
data_collator = DataCollatorWithPadding(tokenizer=tokenizer)

# 3. 加载模型
model = AutoModelForSequenceClassification.from_pretrained(
    "bert-base-uncased", 
    num_labels=2
)

# 4. 定义评估指标
metric = evaluate.load("accuracy")
def compute_metrics(eval_pred):
    logits, labels = eval_pred
    predictions = torch.argmax(logits, dim=-1)
    return metric.compute(predictions=predictions, references=labels)

# 5. 配置训练参数
training_args = TrainingArguments(
    output_dir="./results",
    
    learning_rate=2e-5,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=32,
    num_train_epochs=3,
    weight_decay=0.01,
    
    load_best_model_at_end=False
)

# 6. 创建Trainer实例
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_data["train"].select(range(1000)),
    eval_dataset=tokenized_data["test"].select(range(200)),
    data_collator=data_collator,
    compute_metrics=compute_metrics
)

# 7. 开始训练
trainer.train()